package com.suxin.kafka.thread;

import lombok.Data;
import lombok.extern.slf4j.Slf4j;

import java.util.Objects;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author Tang
 * @classname KafkaConsumerPool
 * @description [ 多綫程顺序消费 ]
 * @date 2022/2/15 21:05
 */
@Slf4j
@Data
public class KafkaConsumerPool<E> {

    /**
     * 线程并发等级
     */
    private Integer concurrentSize;

    /**
     * 工作的线程数
     */
    private Thread[] workThreads;

    /**
     * 任务处理队列
     */
    private LinkedBlockingQueue<E>[] queues;

    private AtomicInteger atomicInteger = new AtomicInteger(0);

    private KafkaConsumerConfig<E> consumerConfig;

    private Runnable runnable = () -> {
        while (true) {
            E task = poll(this.atomicInteger.get());
            if (Objects.nonNull(task)) {
                this.consumerConfig.getBizService().accept(task);
            }
        }
    };

    public KafkaConsumerPool(KafkaConsumerConfig<E> consumerConfig) {
        this.consumerConfig = consumerConfig;
        this.concurrentSize = consumerConfig.getConcurrentSize();
        this.workThreads = new Thread[this.concurrentSize];
        this.queues = new LinkedBlockingQueue[this.concurrentSize];
        this.initWork();
    }

    private void initWork() {
        for (int i = 0; i < this.getConcurrentSize(); i++) {
            Thread thread = new Thread(runnable, this.consumerConfig.getBizName() + "-" + atomicInteger.incrementAndGet());
            this.workThreads[i] = thread;
        }
        for (int i = 0; i < this.getConcurrentSize(); i++) {
            this.workThreads[i].start();
        }
    }

    public void addTask(int index, E e) {
        try {
            this.queues[index % this.getConcurrentSize()].put(e);
        } catch (InterruptedException interruptedException) {
            interruptedException.printStackTrace();
        }
    }

    private E poll(int i) {
        try {
            return this.queues[i].take();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        return null;
    }

}